#include "atomic_0701_ddzt.h"

// 初始化扩展卡尔曼滤波器
void initEKF(EKFState *ekf)
{
    for (int i = 0; i < 6; i++)
    {
        ekf->x[i] = 0.0;
        for (int j = 0; j < 6; j++)
        {
            ekf->P[i][j] = (i == j) ? 1.0 : 0.0;
            ekf->Q[i][j] = (i == j) ? 0.1 : 0.0;
        }
    }

    for (int i = 0; i < 3; i++)
    {
        for (int j = 0; j < 3; j++)
        {
            ekf->R[i][j] = (i == j) ? 0.5 : 0.0;
        }
    }
}

// 扩展卡尔曼滤波器预测步骤
void predictEKF(EKFState *ekf, double dt, double ax, double ay, double az)
{
    // 状态转移矩阵F
    double F[6][6] = {0};
    for (int i = 0; i < 6; i++)
        F[i][i] = 1.0;
    F[0][3] = dt;
    F[1][4] = dt;
    F[2][5] = dt;

    // 控制输入矩阵B
    double B[6][3] = {0};
    B[3][0] = dt;
    B[4][1] = dt;
    B[5][2] = dt;

    // 预测状态
    double x_pred[6] = {0};
    for (int i = 0; i < 6; i++)
    {
        x_pred[i] = ekf->x[i];
        for (int j = 0; j < 6; j++)
        {
            x_pred[i] += F[i][j] * ekf->x[j];
        }
        if (i >= 3)
        {
            x_pred[i] += B[i][i - 3] * (i == 3 ? ax : (i == 4 ? ay : az));
        }
    }

    // 预测协方差
    double P_pred[6][6] = {0};
    for (int i = 0; i < 6; i++)
    {
        for (int j = 0; j < 6; j++)
        {
            P_pred[i][j] = 0.0;
            for (int k = 0; k < 6; k++)
            {
                P_pred[i][j] += F[i][k] * ekf->P[k][j];
            }
            for (int k = 0; k < 6; k++)
            {
                P_pred[i][j] += ekf->P[i][k] * F[j][k];
            }
            P_pred[i][j] += ekf->Q[i][j];
        }
    }

    // 更新状态和协方差
    for (int i = 0; i < 6; i++)
    {
        ekf->x[i] = x_pred[i];
        for (int j = 0; j < 6; j++)
        {
            ekf->P[i][j] = P_pred[i][j];
        }
    }
}

// 扩展卡尔曼滤波器更新步骤
void updateEKF(EKFState *ekf, double zx, double zy, double zz)
{
    // 测量矩阵H
    double H[3][6] = {0};
    H[0][0] = 1.0;
    H[1][1] = 1.0;
    H[2][2] = 1.0;

    // 计算卡尔曼增益
    double S[3][3] = {0};
    for (int i = 0; i < 3; i++)
    {
        for (int j = 0; j < 3; j++)
        {
            S[i][j] = ekf->R[i][j];
            for (int k = 0; k < 6; k++)
            {
                S[i][j] += H[i][k] * ekf->P[k][j] * H[j][k];
            }
        }
    }

    double K[6][3] = {0};
    for (int i = 0; i < 6; i++)
    {
        for (int j = 0; j < 3; j++)
        {
            for (int k = 0; k < 6; k++)
            {
                K[i][j] += ekf->P[i][k] * H[j][k];
            }
            K[i][j] /= S[j][j];
        }
    }

    // 更新状态
    double z[3] = {zx, zy, zz};
    double y[3] = {0};
    for (int i = 0; i < 3; i++)
    {
        y[i] = z[i];
        for (int j = 0; j < 6; j++)
        {
            y[i] -= H[i][j] * ekf->x[j];
        }
    }

    for (int i = 0; i < 6; i++)
    {
        for (int j = 0; j < 3; j++)
        {
            ekf->x[i] += K[i][j] * y[j];
        }
    }

    // 更新协方差
    double I[6][6] = {0};
    for (int i = 0; i < 6; i++)
        I[i][i] = 1.0;

    double P_temp[6][6] = {0};
    for (int i = 0; i < 6; i++)
    {
        for (int j = 0; j < 6; j++)
        {
            P_temp[i][j] = 0.0;
            for (int k = 0; k < 3; k++)
            {
                P_temp[i][j] -= K[i][k] * H[k][j];
            }
            P_temp[i][j] += I[i][j];
        }
    }

    double P_new[6][6] = {0};
    for (int i = 0; i < 6; i++)
    {
        for (int j = 0; j < 6; j++)
        {
            P_new[i][j] = 0.0;
            for (int k = 0; k < 6; k++)
            {
                P_new[i][j] += P_temp[i][k] * ekf->P[k][j];
            }
        }
    }

    for (int i = 0; i < 6; i++)
    {
        for (int j = 0; j < 6; j++)
        {
            ekf->P[i][j] = P_new[i][j];
        }
    }
}

// 导弹状态函数
int getMissileStateData(MissileStateInput *input, MissileStateOutput *output)
{
    static EKFState ekf;
    static int is_initialized = 0;
    static double prev_velocity = 0.0;
    static double prev_altitude = 0.0;
    static double prev_angle = 0.0;
    static double prev_time = 0.0;

    if (!is_initialized)
    {
        initEKF(&ekf);
        prev_velocity = input->velocity;
        prev_altitude = input->altitude;
        prev_angle = input->angle;
        prev_time = input->timestamp;
        is_initialized = 1;
    }

    double dt = input->timestamp - prev_time;
    if (dt <= 0.0)
        dt = 0.01; // 最小时间间隔

    // 扩展卡尔曼滤波
    predictEKF(&ekf, dt, input->acceleration_x, input->acceleration_y, input->acceleration_z);
    updateEKF(&ekf, input->gps_latitude, input->gps_longitude, input->altitude);

    // 计算速度、高度和角度
    double velocity = sqrt(ekf.x[3] * ekf.x[3] + ekf.x[4] * ekf.x[4] + ekf.x[5] * ekf.x[5]);
    double altitude = ekf.x[2];
    double angle = atan2(ekf.x[4], ekf.x[3]) * 180.0 / PI;

    // 计算角速度
    double yaw_rate = (angle - prev_angle) / dt;
    double pitch_rate = (input->angle - prev_angle) / dt;
    double roll_rate = sqrt(input->acceleration_y * input->acceleration_y + input->acceleration_z * input->acceleration_z) / 9.81;

    // 低通滤波加速度
    static double filtered_acc_x = 0.0, filtered_acc_y = 0.0, filtered_acc_z = 0.0;
    double alpha = 0.2;
    filtered_acc_x = alpha * input->acceleration_x + (1 - alpha) * filtered_acc_x;
    filtered_acc_y = alpha * input->acceleration_y + (1 - alpha) * filtered_acc_y;
    filtered_acc_z = alpha * input->acceleration_z + (1 - alpha) * filtered_acc_z;

    // 计算误差
    double velocity_error = velocity - prev_velocity;
    double altitude_error = altitude - prev_altitude;
    double angle_error = angle - prev_angle;

    // 更新状态
    prev_velocity = velocity;
    prev_altitude = altitude;
    prev_angle = angle;
    prev_time = input->timestamp;

    // 填充输出结构体
    output->missile_velocity = velocity;
    output->missile_altitude = altitude;
    output->missile_angle = angle;
    output->yaw_rate = yaw_rate;
    output->pitch_rate = pitch_rate;
    output->roll_rate = roll_rate;
    output->filtered_acceleration_x = filtered_acc_x;
    output->filtered_acceleration_y = filtered_acc_y;
    output->filtered_acceleration_z = filtered_acc_z;
    output->estimated_latitude = ekf.x[0];
    output->estimated_longitude = ekf.x[1];
    output->velocity_error = velocity_error;
    output->altitude_error = altitude_error;
    output->angle_error = angle_error;
    output->kalman_gain_velocity = ekf.P[3][3];
    output->kalman_gain_altitude = ekf.P[2][2];
    output->kalman_gain_angle = atan2(ekf.P[4][4], ekf.P[3][3]) * 180.0 / PI;
    output->covariance_velocity = ekf.P[3][3] + ekf.P[4][4] + ekf.P[5][5];
    output->covariance_altitude = ekf.P[2][2];
    output->covariance_angle = atan2(ekf.P[4][4], ekf.P[3][3]) * 180.0 / PI;

    // 添加一些复杂的计算和模拟
    for (int i = 0; i < 100; i++)
    {
        // 模拟一些复杂的矩阵运算
        double temp = output->missile_velocity * sin(output->missile_angle * PI / 180.0);
        output->missile_altitude += temp * dt / 100.0;

        // 模拟一些非线性效应
        double drag = 0.5 * 1.225 * pow(output->missile_velocity, 2) * 0.1;
        output->missile_velocity -= drag * dt / 100.0;

        // 模拟科里奥利力效应
        double coriolis = 2 * 7.2921159e-5 * output->missile_velocity * sin(input->gps_latitude * PI / 180.0);
        output->yaw_rate += coriolis * dt / 100.0;
    }

    return ATOMIC_SERVICE_SUCCESS;
}